---

title: TSTPlus (Time Series Transformer)


keywords: fastai
sidebar: home_sidebar

summary: "This is an unofficial PyTorch implementation created by Ignacio Oguiza (timeseriesAI@gmail.com) based on TST (Zerveas, 2020) and Transformer (Vaswani, 2017)."
description: "This is an unofficial PyTorch implementation created by Ignacio Oguiza (timeseriesAI@gmail.com) based on TST (Zerveas, 2020) and Transformer (Vaswani, 2017)."
nb_path: "nbs/108c_models.TSTPlus.ipynb"
---
<!--

#################################################
### THIS FILE WAS AUTOGENERATED! DO NOT EDIT! ###
#################################################
# file to edit: nbs/108c_models.TSTPlus.ipynb
# command to build the docs after a change: nbdev_build_docs

-->

<div class="container" id="notebook-container">
        
    {% raw %}
    
<div class="cell border-box-sizing code_cell rendered">

</div>
    {% endraw %}

<div class="cell border-box-sizing text_cell rendered"><div class="inner_cell">
<div class="text_cell_render border-box-sizing rendered_html">
<p><strong>References:</strong></p>
<p>This is an unofficial PyTorch implementation by Ignacio Oguiza of  - oguiza@gmail.com based on:</p>
<ul>
<li>George Zerveas et al. A Transformer-based Framework for Multivariate Time Series Representation Learning, in Proceedings of the 27th ACM SIGKDD Conference on Knowledge Discovery and Data Mining (KDD '21), August 14--18, 2021. ArXiV version: <a href="https://arxiv.org/abs/2010.02803">https://arxiv.org/abs/2010.02803</a></li>
<li>Official implementation: <a href="https://github.com/gzerveas/mvts_transformer">https://github.com/gzerveas/mvts_transformer</a></li>
</ul>
<div class="highlight"><pre><span></span>@inproceedings<span class="o">{</span><span class="m">10</span>.1145/3447548.3467401,
<span class="nv">author</span> <span class="o">=</span> <span class="o">{</span>Zerveas, George and Jayaraman, Srideepika and Patel, Dhaval and Bhamidipaty, Anuradha and Eickhoff, Carsten<span class="o">}</span>,
<span class="nv">title</span> <span class="o">=</span> <span class="o">{</span>A Transformer-Based Framework <span class="k">for</span> Multivariate Time Series Representation Learning<span class="o">}</span>,
<span class="nv">year</span> <span class="o">=</span> <span class="o">{</span><span class="m">2021</span><span class="o">}</span>,
<span class="nv">isbn</span> <span class="o">=</span> <span class="o">{</span><span class="m">9781450383325</span><span class="o">}</span>,
<span class="nv">publisher</span> <span class="o">=</span> <span class="o">{</span>Association <span class="k">for</span> Computing Machinery<span class="o">}</span>,
<span class="nv">address</span> <span class="o">=</span> <span class="o">{</span>New York, NY, USA<span class="o">}</span>,
<span class="nv">url</span> <span class="o">=</span> <span class="o">{</span>https://doi.org/10.1145/3447548.3467401<span class="o">}</span>,
<span class="nv">doi</span> <span class="o">=</span> <span class="o">{</span><span class="m">10</span>.1145/3447548.3467401<span class="o">}</span>,
<span class="nv">booktitle</span> <span class="o">=</span> <span class="o">{</span>Proceedings of the 27th ACM SIGKDD Conference on Knowledge Discovery <span class="p">&amp;</span>amp<span class="p">;</span> Data Mining<span class="o">}</span>,
<span class="nv">pages</span> <span class="o">=</span> <span class="o">{</span><span class="m">2114</span>–2124<span class="o">}</span>,
<span class="nv">numpages</span> <span class="o">=</span> <span class="o">{</span><span class="m">11</span><span class="o">}</span>,
<span class="nv">keywords</span> <span class="o">=</span> <span class="o">{</span>regression, framework, multivariate <span class="nb">time</span> series, classification, transformer, deep learning, self-supervised learning, unsupervised learning, imputation<span class="o">}</span>,
<span class="nv">location</span> <span class="o">=</span> <span class="o">{</span>Virtual Event, Singapore<span class="o">}</span>,
<span class="nv">series</span> <span class="o">=</span> <span class="o">{</span>KDD <span class="err">&#39;</span><span class="m">21</span><span class="o">}</span>
<span class="o">}</span>
</pre></div>
<ul>
<li><p>Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., ... &amp; Polosukhin, I. (2017). <a href="https://papers.nips.cc/paper/7181-attention-is-all-you-need.pdf">Attention is all you need. In Advances in neural information processing systems (pp. 5998-6008).</a></p>
</li>
<li><p>He, R., Ravula, A., Kanagal, B., &amp; Ainslie, J. (2020). Realformer: Transformer Likes Informed Attention. arXiv preprint arXiv:2012.11747.</p>
</li>
</ul>
<p>This implementation is adapted to work with the rest of the <code>tsai</code> library, and contain some hyperparameters that are not available in the original implementation. I included them for experimenting.</p>

</div>
</div>
</div>
<div class="cell border-box-sizing text_cell rendered"><div class="inner_cell">
<div class="text_cell_render border-box-sizing rendered_html">
<h2 id="Imports">Imports<a class="anchor-link" href="#Imports"> </a></h2>
</div>
</div>
</div>
    {% raw %}
    
<div class="cell border-box-sizing code_cell rendered">

</div>
    {% endraw %}

<div class="cell border-box-sizing text_cell rendered"><div class="inner_cell">
<div class="text_cell_render border-box-sizing rendered_html">
<h2 id="Positional-encoders">Positional encoders<a class="anchor-link" href="#Positional-encoders"> </a></h2>
</div>
</div>
</div>
    {% raw %}
    
<div class="cell border-box-sizing code_cell rendered">

<div class="output_wrapper">
<div class="output">

<div class="output_area">


<div class="output_markdown rendered_html output_subarea ">
<h4 id="PositionalEncoding" class="doc_header"><code>PositionalEncoding</code><a href="https://github.com/timeseriesAI/tsai/tree/main/tsai/models/TSTPlus.py#L14" class="source_link" style="float:right">[source]</a></h4><blockquote><p><code>PositionalEncoding</code>(<strong><code>q_len</code></strong>, <strong><code>d_model</code></strong>, <strong><code>normalize</code></strong>=<em><code>True</code></em>)</p>
</blockquote>

</div>

</div>

</div>
</div>

</div>
    {% endraw %}

    {% raw %}
    
<div class="cell border-box-sizing code_cell rendered">

</div>
    {% endraw %}

    {% raw %}
    
<div class="cell border-box-sizing code_cell rendered">
<div class="input">

<div class="inner_cell">
    <div class="input_area">
<div class=" highlight hl-ipython3"><pre><span></span><span class="n">pe</span> <span class="o">=</span> <span class="n">PositionalEncoding</span><span class="p">(</span><span class="mi">1000</span><span class="p">,</span> <span class="mi">512</span><span class="p">)</span><span class="o">.</span><span class="n">detach</span><span class="p">()</span><span class="o">.</span><span class="n">cpu</span><span class="p">()</span><span class="o">.</span><span class="n">numpy</span><span class="p">()</span>
<span class="n">plt</span><span class="o">.</span><span class="n">pcolormesh</span><span class="p">(</span><span class="n">pe</span><span class="p">,</span> <span class="n">cmap</span><span class="o">=</span><span class="s1">&#39;viridis&#39;</span><span class="p">)</span>
<span class="n">plt</span><span class="o">.</span><span class="n">title</span><span class="p">(</span><span class="s1">&#39;PositionalEncoding&#39;</span><span class="p">)</span>
<span class="n">plt</span><span class="o">.</span><span class="n">colorbar</span><span class="p">()</span>
<span class="n">plt</span><span class="o">.</span><span class="n">show</span><span class="p">()</span>
<span class="n">pe</span><span class="o">.</span><span class="n">mean</span><span class="p">(),</span> <span class="n">pe</span><span class="o">.</span><span class="n">std</span><span class="p">(),</span> <span class="n">pe</span><span class="o">.</span><span class="n">min</span><span class="p">(),</span> <span class="n">pe</span><span class="o">.</span><span class="n">max</span><span class="p">(),</span> <span class="n">pe</span><span class="o">.</span><span class="n">shape</span>
</pre></div>

    </div>
</div>
</div>

<div class="output_wrapper">
<div class="output">

<div class="output_area">



<div class="output_png output_subarea ">
<img src="
"
>
</div>

</div>

<div class="output_area">



<div class="output_text output_subarea output_execute_result">
<pre>(3.2037498e-10, 0.09999991, -0.18388666, 0.11518021, (1000, 512))</pre>
</div>

</div>

</div>
</div>

</div>
    {% endraw %}

    {% raw %}
    
<div class="cell border-box-sizing code_cell rendered">

<div class="output_wrapper">
<div class="output">

<div class="output_area">


<div class="output_markdown rendered_html output_subarea ">
<h4 id="Coord2dPosEncoding" class="doc_header"><code>Coord2dPosEncoding</code><a href="https://github.com/timeseriesAI/tsai/tree/main/tsai/models/TSTPlus.py#L28" class="source_link" style="float:right">[source]</a></h4><blockquote><p><code>Coord2dPosEncoding</code>(<strong><code>q_len</code></strong>, <strong><code>d_model</code></strong>, <strong><code>exponential</code></strong>=<em><code>False</code></em>, <strong><code>normalize</code></strong>=<em><code>True</code></em>, <strong><code>eps</code></strong>=<em><code>0.001</code></em>, <strong><code>verbose</code></strong>=<em><code>False</code></em>, <strong><code>device</code></strong>=<em><code>device(type='cpu')</code></em>)</p>
</blockquote>

</div>

</div>

</div>
</div>

</div>
    {% endraw %}

    {% raw %}
    
<div class="cell border-box-sizing code_cell rendered">

</div>
    {% endraw %}

    {% raw %}
    
<div class="cell border-box-sizing code_cell rendered">
<div class="input">

<div class="inner_cell">
    <div class="input_area">
<div class=" highlight hl-ipython3"><pre><span></span><span class="n">cpe</span> <span class="o">=</span> <span class="n">Coord2dPosEncoding</span><span class="p">(</span><span class="mi">1000</span><span class="p">,</span> <span class="mi">512</span><span class="p">,</span> <span class="n">exponential</span><span class="o">=</span><span class="kc">True</span><span class="p">,</span> <span class="n">normalize</span><span class="o">=</span><span class="kc">True</span><span class="p">)</span><span class="o">.</span><span class="n">cpu</span><span class="p">()</span><span class="o">.</span><span class="n">numpy</span><span class="p">()</span>
<span class="n">plt</span><span class="o">.</span><span class="n">pcolormesh</span><span class="p">(</span><span class="n">cpe</span><span class="p">,</span> <span class="n">cmap</span><span class="o">=</span><span class="s1">&#39;viridis&#39;</span><span class="p">)</span>
<span class="n">plt</span><span class="o">.</span><span class="n">title</span><span class="p">(</span><span class="s1">&#39;Coord2dPosEncoding&#39;</span><span class="p">)</span>
<span class="n">plt</span><span class="o">.</span><span class="n">colorbar</span><span class="p">()</span>
<span class="n">plt</span><span class="o">.</span><span class="n">show</span><span class="p">()</span>
<span class="n">plt</span><span class="o">.</span><span class="n">plot</span><span class="p">(</span><span class="n">cpe</span><span class="o">.</span><span class="n">mean</span><span class="p">(</span><span class="mi">0</span><span class="p">))</span>
<span class="n">plt</span><span class="o">.</span><span class="n">show</span><span class="p">()</span>
<span class="n">plt</span><span class="o">.</span><span class="n">plot</span><span class="p">(</span><span class="n">cpe</span><span class="o">.</span><span class="n">mean</span><span class="p">(</span><span class="mi">1</span><span class="p">))</span>
<span class="n">plt</span><span class="o">.</span><span class="n">show</span><span class="p">()</span>
<span class="n">cpe</span><span class="o">.</span><span class="n">mean</span><span class="p">(),</span> <span class="n">cpe</span><span class="o">.</span><span class="n">std</span><span class="p">(),</span> <span class="n">cpe</span><span class="o">.</span><span class="n">min</span><span class="p">(),</span> <span class="n">cpe</span><span class="o">.</span><span class="n">max</span><span class="p">()</span>
</pre></div>

    </div>
</div>
</div>

<div class="output_wrapper">
<div class="output">

<div class="output_area">



<div class="output_png output_subarea ">
<img src="
"
>
</div>

</div>

<div class="output_area">



<div class="output_png output_subarea ">
<img src="
"
>
</div>

</div>

<div class="output_area">



<div class="output_png output_subarea ">
<img src="
"
>
</div>

</div>

<div class="output_area">



<div class="output_text output_subarea output_execute_result">
<pre>(3.695488e-09, 0.09999991, -0.22459325, 0.22487777)</pre>
</div>

</div>

</div>
</div>

</div>
    {% endraw %}

    {% raw %}
    
<div class="cell border-box-sizing code_cell rendered">

<div class="output_wrapper">
<div class="output">

<div class="output_area">


<div class="output_markdown rendered_html output_subarea ">
<h4 id="Coord1dPosEncoding" class="doc_header"><code>Coord1dPosEncoding</code><a href="https://github.com/timeseriesAI/tsai/tree/main/tsai/models/TSTPlus.py#L44" class="source_link" style="float:right">[source]</a></h4><blockquote><p><code>Coord1dPosEncoding</code>(<strong><code>q_len</code></strong>, <strong><code>exponential</code></strong>=<em><code>False</code></em>, <strong><code>normalize</code></strong>=<em><code>True</code></em>, <strong><code>device</code></strong>=<em><code>device(type='cpu')</code></em>)</p>
</blockquote>

</div>

</div>

</div>
</div>

</div>
    {% endraw %}

    {% raw %}
    
<div class="cell border-box-sizing code_cell rendered">

</div>
    {% endraw %}

    {% raw %}
    
<div class="cell border-box-sizing code_cell rendered">
<div class="input">

<div class="inner_cell">
    <div class="input_area">
<div class=" highlight hl-ipython3"><pre><span></span><span class="n">cpe</span> <span class="o">=</span> <span class="n">Coord1dPosEncoding</span><span class="p">(</span><span class="mi">1000</span><span class="p">,</span> <span class="n">exponential</span><span class="o">=</span><span class="kc">True</span><span class="p">,</span> <span class="n">normalize</span><span class="o">=</span><span class="kc">True</span><span class="p">)</span><span class="o">.</span><span class="n">detach</span><span class="p">()</span><span class="o">.</span><span class="n">cpu</span><span class="p">()</span><span class="o">.</span><span class="n">numpy</span><span class="p">()</span>
<span class="n">plt</span><span class="o">.</span><span class="n">pcolormesh</span><span class="p">(</span><span class="n">cpe</span><span class="p">,</span> <span class="n">cmap</span><span class="o">=</span><span class="s1">&#39;viridis&#39;</span><span class="p">)</span>
<span class="n">plt</span><span class="o">.</span><span class="n">title</span><span class="p">(</span><span class="s1">&#39;Coord1dPosEncoding&#39;</span><span class="p">)</span>
<span class="n">plt</span><span class="o">.</span><span class="n">colorbar</span><span class="p">()</span>
<span class="n">plt</span><span class="o">.</span><span class="n">show</span><span class="p">()</span>
<span class="n">plt</span><span class="o">.</span><span class="n">plot</span><span class="p">(</span><span class="n">cpe</span><span class="o">.</span><span class="n">mean</span><span class="p">(</span><span class="mi">1</span><span class="p">))</span>
<span class="n">plt</span><span class="o">.</span><span class="n">show</span><span class="p">()</span>
<span class="n">cpe</span><span class="o">.</span><span class="n">mean</span><span class="p">(),</span> <span class="n">cpe</span><span class="o">.</span><span class="n">std</span><span class="p">(),</span> <span class="n">cpe</span><span class="o">.</span><span class="n">min</span><span class="p">(),</span> <span class="n">cpe</span><span class="o">.</span><span class="n">max</span><span class="p">(),</span> <span class="n">cpe</span><span class="o">.</span><span class="n">shape</span>
</pre></div>

    </div>
</div>
</div>

<div class="output_wrapper">
<div class="output">

<div class="output_area">



<div class="output_png output_subarea ">
<img src="
"
>
</div>

</div>

<div class="output_area">



<div class="output_png output_subarea ">
<img src="
"
>
</div>

</div>

<div class="output_area">



<div class="output_text output_subarea output_execute_result">
<pre>(0.0, 0.099949986, -0.2820423, 0.14113107, (1000, 1))</pre>
</div>

</div>

</div>
</div>

</div>
    {% endraw %}

    {% raw %}
    
<div class="cell border-box-sizing code_cell rendered">
<div class="input">

<div class="inner_cell">
    <div class="input_area">
<div class=" highlight hl-ipython3"><pre><span></span><span class="n">cpe</span> <span class="o">=</span> <span class="n">Coord1dPosEncoding</span><span class="p">(</span><span class="mi">1000</span><span class="p">,</span> <span class="n">exponential</span><span class="o">=</span><span class="kc">True</span><span class="p">,</span> <span class="n">normalize</span><span class="o">=</span><span class="kc">True</span><span class="p">)</span><span class="o">.</span><span class="n">detach</span><span class="p">()</span><span class="o">.</span><span class="n">cpu</span><span class="p">()</span><span class="o">.</span><span class="n">numpy</span><span class="p">()</span>
<span class="n">plt</span><span class="o">.</span><span class="n">pcolormesh</span><span class="p">(</span><span class="n">cpe</span><span class="p">,</span> <span class="n">cmap</span><span class="o">=</span><span class="s1">&#39;viridis&#39;</span><span class="p">)</span>
<span class="n">plt</span><span class="o">.</span><span class="n">title</span><span class="p">(</span><span class="s1">&#39;Coord1dPosEncoding&#39;</span><span class="p">)</span>
<span class="n">plt</span><span class="o">.</span><span class="n">colorbar</span><span class="p">()</span>
<span class="n">plt</span><span class="o">.</span><span class="n">show</span><span class="p">()</span>
<span class="n">plt</span><span class="o">.</span><span class="n">plot</span><span class="p">(</span><span class="n">cpe</span><span class="o">.</span><span class="n">mean</span><span class="p">(</span><span class="mi">1</span><span class="p">))</span>
<span class="n">plt</span><span class="o">.</span><span class="n">show</span><span class="p">()</span>
<span class="n">cpe</span><span class="o">.</span><span class="n">mean</span><span class="p">(),</span> <span class="n">cpe</span><span class="o">.</span><span class="n">std</span><span class="p">(),</span> <span class="n">cpe</span><span class="o">.</span><span class="n">min</span><span class="p">(),</span> <span class="n">cpe</span><span class="o">.</span><span class="n">max</span><span class="p">()</span>
</pre></div>

    </div>
</div>
</div>

<div class="output_wrapper">
<div class="output">

<div class="output_area">



<div class="output_png output_subarea ">
<img src="
"
>
</div>

</div>

<div class="output_area">



<div class="output_png output_subarea ">
<img src="
"
>
</div>

</div>

<div class="output_area">



<div class="output_text output_subarea output_execute_result">
<pre>(0.0, 0.099949986, -0.2820423, 0.14113107)</pre>
</div>

</div>

</div>
</div>

</div>
    {% endraw %}

<div class="cell border-box-sizing text_cell rendered"><div class="inner_cell">
<div class="text_cell_render border-box-sizing rendered_html">
<h2 id="TST">TST<a class="anchor-link" href="#TST"> </a></h2>
</div>
</div>
</div>
    {% raw %}
    
<div class="cell border-box-sizing code_cell rendered">
<div class="input">

<div class="inner_cell">
    <div class="input_area">
<div class=" highlight hl-ipython3"><pre><span></span><span class="n">t</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">rand</span><span class="p">(</span><span class="mi">16</span><span class="p">,</span> <span class="mi">50</span><span class="p">,</span> <span class="mi">128</span><span class="p">)</span>
<span class="n">attn_mask</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">triu</span><span class="p">(</span><span class="n">torch</span><span class="o">.</span><span class="n">ones</span><span class="p">(</span><span class="mi">50</span><span class="p">,</span> <span class="mi">50</span><span class="p">))</span> <span class="c1"># shape: q_len x q_len</span>
<span class="n">key_padding_mask</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">zeros</span><span class="p">(</span><span class="mi">16</span><span class="p">,</span> <span class="mi">50</span><span class="p">)</span>
<span class="n">key_padding_mask</span><span class="p">[[</span><span class="mi">1</span><span class="p">,</span> <span class="mi">3</span><span class="p">,</span> <span class="mi">6</span><span class="p">,</span> <span class="mi">15</span><span class="p">],</span> <span class="o">-</span><span class="mi">10</span><span class="p">:]</span> <span class="o">=</span> <span class="mi">1</span>
<span class="n">key_padding_mask</span> <span class="o">=</span> <span class="n">key_padding_mask</span><span class="o">.</span><span class="n">bool</span><span class="p">()</span>
<span class="nb">print</span><span class="p">(</span><span class="s1">&#39;attn_mask&#39;</span><span class="p">,</span> <span class="n">attn_mask</span><span class="o">.</span><span class="n">shape</span><span class="p">,</span> <span class="s1">&#39;key_padding_mask&#39;</span><span class="p">,</span> <span class="n">key_padding_mask</span><span class="o">.</span><span class="n">shape</span><span class="p">)</span>
<span class="n">encoder</span> <span class="o">=</span> <span class="n">_TSTEncoderLayer</span><span class="p">(</span><span class="n">q_len</span><span class="o">=</span><span class="mi">50</span><span class="p">,</span> <span class="n">d_model</span><span class="o">=</span><span class="mi">128</span><span class="p">,</span> <span class="n">n_heads</span><span class="o">=</span><span class="mi">8</span><span class="p">,</span> <span class="n">d_k</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">d_v</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">d_ff</span><span class="o">=</span><span class="mi">512</span><span class="p">,</span> <span class="n">attn_dropout</span><span class="o">=</span><span class="mf">0.</span><span class="p">,</span> <span class="n">dropout</span><span class="o">=</span><span class="mf">0.1</span><span class="p">,</span> <span class="n">store_attn</span><span class="o">=</span><span class="kc">True</span><span class="p">,</span> <span class="n">activation</span><span class="o">=</span><span class="s1">&#39;gelu&#39;</span><span class="p">)</span>
<span class="n">output</span> <span class="o">=</span> <span class="n">encoder</span><span class="p">(</span><span class="n">t</span><span class="p">,</span> <span class="n">key_padding_mask</span><span class="o">=</span><span class="n">key_padding_mask</span><span class="p">,</span> <span class="n">attn_mask</span><span class="o">=</span><span class="n">attn_mask</span><span class="p">)</span>
<span class="n">output</span><span class="o">.</span><span class="n">shape</span>
</pre></div>

    </div>
</div>
</div>

<div class="output_wrapper">
<div class="output">

<div class="output_area">

<div class="output_subarea output_stream output_stdout output_text">
<pre>attn_mask torch.Size([50, 50]) key_padding_mask torch.Size([16, 50])
</pre>
</div>
</div>

<div class="output_area">



<div class="output_text output_subarea output_execute_result">
<pre>torch.Size([16, 50, 128])</pre>
</div>

</div>

</div>
</div>

</div>
    {% endraw %}

    {% raw %}
    
<div class="cell border-box-sizing code_cell rendered">
<div class="input">

<div class="inner_cell">
    <div class="input_area">
<div class=" highlight hl-ipython3"><pre><span></span><span class="n">cmap</span><span class="o">=</span><span class="s1">&#39;viridis&#39;</span>
<span class="n">figsize</span><span class="o">=</span><span class="p">(</span><span class="mi">6</span><span class="p">,</span><span class="mi">5</span><span class="p">)</span>
<span class="n">plt</span><span class="o">.</span><span class="n">figure</span><span class="p">(</span><span class="n">figsize</span><span class="o">=</span><span class="n">figsize</span><span class="p">)</span>
<span class="n">plt</span><span class="o">.</span><span class="n">pcolormesh</span><span class="p">(</span><span class="n">encoder</span><span class="o">.</span><span class="n">attn</span><span class="p">[</span><span class="mi">0</span><span class="p">][</span><span class="mi">0</span><span class="p">]</span><span class="o">.</span><span class="n">detach</span><span class="p">()</span><span class="o">.</span><span class="n">cpu</span><span class="p">()</span><span class="o">.</span><span class="n">numpy</span><span class="p">(),</span> <span class="n">cmap</span><span class="o">=</span><span class="n">cmap</span><span class="p">)</span>
<span class="n">plt</span><span class="o">.</span><span class="n">title</span><span class="p">(</span><span class="s1">&#39;Self-attention map&#39;</span><span class="p">)</span>
<span class="n">plt</span><span class="o">.</span><span class="n">colorbar</span><span class="p">()</span>
<span class="n">plt</span><span class="o">.</span><span class="n">show</span><span class="p">()</span>
</pre></div>

    </div>
</div>
</div>

<div class="output_wrapper">
<div class="output">

<div class="output_area">



<div class="output_png output_subarea ">
<img src="
"
>
</div>

</div>

</div>
</div>

</div>
    {% endraw %}

    {% raw %}
    
<div class="cell border-box-sizing code_cell rendered">

<div class="output_wrapper">
<div class="output">

<div class="output_area">


<div class="output_markdown rendered_html output_subarea ">
<h2 id="TSTPlus" class="doc_header"><code>class</code> <code>TSTPlus</code><a href="https://github.com/timeseriesAI/tsai/tree/main/tsai/models/TSTPlus.py#L258" class="source_link" style="float:right">[source]</a></h2><blockquote><p><code>TSTPlus</code>(<strong><code>c_in</code></strong>:<code>int</code>, <strong><code>c_out</code></strong>:<code>int</code>, <strong><code>seq_len</code></strong>:<code>int</code>, <strong><code>max_seq_len</code></strong>:<code>Optional</code>[<code>int</code>]=<em><code>512</code></em>, <strong><code>n_layers</code></strong>:<code>int</code>=<em><code>3</code></em>, <strong><code>d_model</code></strong>:<code>int</code>=<em><code>128</code></em>, <strong><code>n_heads</code></strong>:<code>int</code>=<em><code>16</code></em>, <strong><code>d_k</code></strong>:<code>Optional</code>[<code>int</code>]=<em><code>None</code></em>, <strong><code>d_v</code></strong>:<code>Optional</code>[<code>int</code>]=<em><code>None</code></em>, <strong><code>d_ff</code></strong>:<code>int</code>=<em><code>256</code></em>, <strong><code>norm</code></strong>:<code>str</code>=<em><code>'BatchNorm'</code></em>, <strong><code>attn_dropout</code></strong>:<code>float</code>=<em><code>0.0</code></em>, <strong><code>dropout</code></strong>:<code>float</code>=<em><code>0.0</code></em>, <strong><code>act</code></strong>:<code>str</code>=<em><code>'gelu'</code></em>, <strong><code>key_padding_mask</code></strong>:<code>bool</code>=<em><code>'auto'</code></em>, <strong><code>padding_var</code></strong>:<code>Optional</code>[<code>int</code>]=<em><code>None</code></em>, <strong><code>attn_mask</code></strong>:<code>Optional</code>[<code>Tensor</code>]=<em><code>None</code></em>, <strong><code>res_attention</code></strong>:<code>bool</code>=<em><code>True</code></em>, <strong><code>pre_norm</code></strong>:<code>bool</code>=<em><code>False</code></em>, <strong><code>store_attn</code></strong>:<code>bool</code>=<em><code>False</code></em>, <strong><code>pe</code></strong>:<code>str</code>=<em><code>'zeros'</code></em>, <strong><code>learn_pe</code></strong>:<code>bool</code>=<em><code>True</code></em>, <strong><code>flatten</code></strong>:<code>bool</code>=<em><code>True</code></em>, <strong><code>fc_dropout</code></strong>:<code>float</code>=<em><code>0.0</code></em>, <strong><code>concat_pool</code></strong>:<code>bool</code>=<em><code>False</code></em>, <strong><code>bn</code></strong>:<code>bool</code>=<em><code>False</code></em>, <strong><code>custom_head</code></strong>:<code>Optional</code>=<em><code>None</code></em>, <strong><code>y_range</code></strong>:<code>Optional</code>[<code>tuple</code>]=<em><code>None</code></em>, <strong><code>verbose</code></strong>:<code>bool</code>=<em><code>False</code></em>, <strong>**<code>kwargs</code></strong>) :: <a href="/models.TabFusionTransformer.html#Sequential"><code>Sequential</code></a></p>
</blockquote>
<p>TST (Time Series Transformer) is a Transformer that takes continuous time series as inputs</p>

</div>

</div>

</div>
</div>

</div>
    {% endraw %}

    {% raw %}
    
<div class="cell border-box-sizing code_cell rendered">

</div>
    {% endraw %}

    {% raw %}
    
<div class="cell border-box-sizing code_cell rendered">
<div class="input">

<div class="inner_cell">
    <div class="input_area">
<div class=" highlight hl-ipython3"><pre><span></span><span class="kn">from</span> <span class="nn">tsai.models.utils</span> <span class="kn">import</span> <span class="n">build_ts_model</span>

<span class="n">bs</span> <span class="o">=</span> <span class="mi">8</span>
<span class="n">c_in</span> <span class="o">=</span> <span class="mi">9</span>  <span class="c1"># aka channels, features, variables, dimensions</span>
<span class="n">c_out</span> <span class="o">=</span> <span class="mi">2</span>
<span class="n">seq_len</span> <span class="o">=</span> <span class="mi">1_500</span>

<span class="n">xb</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">randn</span><span class="p">(</span><span class="n">bs</span><span class="p">,</span> <span class="n">c_in</span><span class="p">,</span> <span class="n">seq_len</span><span class="p">)</span><span class="o">.</span><span class="n">to</span><span class="p">(</span><span class="n">device</span><span class="p">)</span>

<span class="c1"># standardize by channel by_var based on the training set</span>
<span class="n">xb</span> <span class="o">=</span> <span class="p">(</span><span class="n">xb</span> <span class="o">-</span> <span class="n">xb</span><span class="o">.</span><span class="n">mean</span><span class="p">((</span><span class="mi">0</span><span class="p">,</span> <span class="mi">2</span><span class="p">),</span> <span class="n">keepdim</span><span class="o">=</span><span class="kc">True</span><span class="p">))</span> <span class="o">/</span> <span class="n">xb</span><span class="o">.</span><span class="n">std</span><span class="p">((</span><span class="mi">0</span><span class="p">,</span> <span class="mi">2</span><span class="p">),</span> <span class="n">keepdim</span><span class="o">=</span><span class="kc">True</span><span class="p">)</span>

<span class="c1"># Settings</span>
<span class="n">max_seq_len</span> <span class="o">=</span> <span class="mi">256</span>
<span class="n">d_model</span> <span class="o">=</span> <span class="mi">128</span>
<span class="n">n_heads</span> <span class="o">=</span> <span class="mi">16</span>
<span class="n">d_k</span> <span class="o">=</span> <span class="n">d_v</span> <span class="o">=</span> <span class="kc">None</span>  <span class="c1"># if None --&gt; d_model // n_heads</span>
<span class="n">d_ff</span> <span class="o">=</span> <span class="mi">256</span>
<span class="n">norm</span> <span class="o">=</span> <span class="s2">&quot;BatchNorm&quot;</span>
<span class="n">dropout</span> <span class="o">=</span> <span class="mf">0.1</span>
<span class="n">activation</span> <span class="o">=</span> <span class="s2">&quot;gelu&quot;</span>
<span class="n">n_layers</span> <span class="o">=</span> <span class="mi">3</span>
<span class="n">fc_dropout</span> <span class="o">=</span> <span class="mf">0.1</span>
<span class="n">pe</span> <span class="o">=</span> <span class="kc">None</span>
<span class="n">learn_pe</span> <span class="o">=</span> <span class="kc">True</span>
<span class="n">kwargs</span> <span class="o">=</span> <span class="p">{}</span>

<span class="n">model</span> <span class="o">=</span> <span class="n">TSTPlus</span><span class="p">(</span><span class="n">c_in</span><span class="p">,</span> <span class="n">c_out</span><span class="p">,</span> <span class="n">seq_len</span><span class="p">,</span> <span class="n">max_seq_len</span><span class="o">=</span><span class="n">max_seq_len</span><span class="p">,</span> <span class="n">d_model</span><span class="o">=</span><span class="n">d_model</span><span class="p">,</span> <span class="n">n_heads</span><span class="o">=</span><span class="n">n_heads</span><span class="p">,</span>
                <span class="n">d_k</span><span class="o">=</span><span class="n">d_k</span><span class="p">,</span> <span class="n">d_v</span><span class="o">=</span><span class="n">d_v</span><span class="p">,</span> <span class="n">d_ff</span><span class="o">=</span><span class="n">d_ff</span><span class="p">,</span> <span class="n">norm</span><span class="o">=</span><span class="n">norm</span><span class="p">,</span> <span class="n">dropout</span><span class="o">=</span><span class="n">dropout</span><span class="p">,</span> <span class="n">activation</span><span class="o">=</span><span class="n">activation</span><span class="p">,</span> <span class="n">n_layers</span><span class="o">=</span><span class="n">n_layers</span><span class="p">,</span>
                <span class="n">fc_dropout</span><span class="o">=</span><span class="n">fc_dropout</span><span class="p">,</span> <span class="n">pe</span><span class="o">=</span><span class="n">pe</span><span class="p">,</span> <span class="n">learn_pe</span><span class="o">=</span><span class="n">learn_pe</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">)</span><span class="o">.</span><span class="n">to</span><span class="p">(</span><span class="n">device</span><span class="p">)</span>
<span class="n">test_eq</span><span class="p">(</span><span class="n">model</span><span class="p">(</span><span class="n">xb</span><span class="p">)</span><span class="o">.</span><span class="n">shape</span><span class="p">,</span> <span class="p">[</span><span class="n">bs</span><span class="p">,</span> <span class="n">c_out</span><span class="p">])</span>
<span class="n">test_eq</span><span class="p">(</span><span class="n">model</span><span class="p">[</span><span class="mi">0</span><span class="p">],</span> <span class="n">model</span><span class="o">.</span><span class="n">backbone</span><span class="p">)</span>
<span class="n">test_eq</span><span class="p">(</span><span class="n">model</span><span class="p">[</span><span class="mi">1</span><span class="p">],</span> <span class="n">model</span><span class="o">.</span><span class="n">head</span><span class="p">)</span>
<span class="n">model2</span> <span class="o">=</span> <span class="n">build_ts_model</span><span class="p">(</span><span class="n">TSTPlus</span><span class="p">,</span> <span class="n">c_in</span><span class="p">,</span> <span class="n">c_out</span><span class="p">,</span> <span class="n">seq_len</span><span class="p">,</span> <span class="n">max_seq_len</span><span class="o">=</span><span class="n">max_seq_len</span><span class="p">,</span> <span class="n">d_model</span><span class="o">=</span><span class="n">d_model</span><span class="p">,</span> <span class="n">n_heads</span><span class="o">=</span><span class="n">n_heads</span><span class="p">,</span>
                           <span class="n">d_k</span><span class="o">=</span><span class="n">d_k</span><span class="p">,</span> <span class="n">d_v</span><span class="o">=</span><span class="n">d_v</span><span class="p">,</span> <span class="n">d_ff</span><span class="o">=</span><span class="n">d_ff</span><span class="p">,</span> <span class="n">norm</span><span class="o">=</span><span class="n">norm</span><span class="p">,</span> <span class="n">dropout</span><span class="o">=</span><span class="n">dropout</span><span class="p">,</span> <span class="n">activation</span><span class="o">=</span><span class="n">activation</span><span class="p">,</span> <span class="n">n_layers</span><span class="o">=</span><span class="n">n_layers</span><span class="p">,</span>
                           <span class="n">fc_dropout</span><span class="o">=</span><span class="n">fc_dropout</span><span class="p">,</span> <span class="n">pe</span><span class="o">=</span><span class="n">pe</span><span class="p">,</span> <span class="n">learn_pe</span><span class="o">=</span><span class="n">learn_pe</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">)</span><span class="o">.</span><span class="n">to</span><span class="p">(</span><span class="n">device</span><span class="p">)</span>
<span class="n">test_eq</span><span class="p">(</span><span class="n">model2</span><span class="p">(</span><span class="n">xb</span><span class="p">)</span><span class="o">.</span><span class="n">shape</span><span class="p">,</span> <span class="p">[</span><span class="n">bs</span><span class="p">,</span> <span class="n">c_out</span><span class="p">])</span>
<span class="n">test_eq</span><span class="p">(</span><span class="n">model2</span><span class="p">[</span><span class="mi">0</span><span class="p">],</span> <span class="n">model2</span><span class="o">.</span><span class="n">backbone</span><span class="p">)</span>
<span class="n">test_eq</span><span class="p">(</span><span class="n">model2</span><span class="p">[</span><span class="mi">1</span><span class="p">],</span> <span class="n">model2</span><span class="o">.</span><span class="n">head</span><span class="p">)</span>
<span class="nb">print</span><span class="p">(</span><span class="sa">f</span><span class="s1">&#39;model parameters: </span><span class="si">{</span><span class="n">count_parameters</span><span class="p">(</span><span class="n">model</span><span class="p">)</span><span class="si">}</span><span class="s1">&#39;</span><span class="p">)</span>
</pre></div>

    </div>
</div>
</div>

<div class="output_wrapper">
<div class="output">

<div class="output_area">

<div class="output_subarea output_stream output_stdout output_text">
<pre>model parameters: 470018
</pre>
</div>
</div>

</div>
</div>

</div>
    {% endraw %}

    {% raw %}
    
<div class="cell border-box-sizing code_cell rendered">
<div class="input">

<div class="inner_cell">
    <div class="input_area">
<div class=" highlight hl-ipython3"><pre><span></span><span class="n">key_padding_mask</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">sort</span><span class="p">(</span><span class="n">torch</span><span class="o">.</span><span class="n">randint</span><span class="p">(</span><span class="mi">0</span><span class="p">,</span> <span class="mi">2</span><span class="p">,</span> <span class="p">(</span><span class="n">bs</span><span class="p">,</span> <span class="n">max_seq_len</span><span class="p">)))</span><span class="o">.</span><span class="n">values</span><span class="o">.</span><span class="n">bool</span><span class="p">()</span><span class="o">.</span><span class="n">to</span><span class="p">(</span><span class="n">device</span><span class="p">)</span>
<span class="n">key_padding_mask</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span>
</pre></div>

    </div>
</div>
</div>

<div class="output_wrapper">
<div class="output">

<div class="output_area">



<div class="output_text output_subarea output_execute_result">
<pre>tensor([False, False, False, False, False, False, False, False, False, False,
        False, False, False, False, False, False, False, False, False, False,
        False, False, False, False, False, False, False, False, False, False,
        False, False, False, False, False, False, False, False, False, False,
        False, False, False, False, False, False, False, False, False, False,
        False, False, False, False, False, False, False, False, False, False,
        False, False, False, False, False, False, False, False, False, False,
        False, False, False, False, False, False, False, False, False, False,
        False, False, False, False, False, False, False, False, False, False,
        False, False, False, False, False, False, False, False, False, False,
        False, False, False, False, False, False, False, False, False, False,
        False, False, False, False, False, False, False,  True,  True,  True,
         True,  True,  True,  True,  True,  True,  True,  True,  True,  True,
         True,  True,  True,  True,  True,  True,  True,  True,  True,  True,
         True,  True,  True,  True,  True,  True,  True,  True,  True,  True,
         True,  True,  True,  True,  True,  True,  True,  True,  True,  True,
         True,  True,  True,  True,  True,  True,  True,  True,  True,  True,
         True,  True,  True,  True,  True,  True,  True,  True,  True,  True,
         True,  True,  True,  True,  True,  True,  True,  True,  True,  True,
         True,  True,  True,  True,  True,  True,  True,  True,  True,  True,
         True,  True,  True,  True,  True,  True,  True,  True,  True,  True,
         True,  True,  True,  True,  True,  True,  True,  True,  True,  True,
         True,  True,  True,  True,  True,  True,  True,  True,  True,  True,
         True,  True,  True,  True,  True,  True,  True,  True,  True,  True,
         True,  True,  True,  True,  True,  True,  True,  True,  True,  True,
         True,  True,  True,  True,  True,  True])</pre>
</div>

</div>

</div>
</div>

</div>
    {% endraw %}

    {% raw %}
    
<div class="cell border-box-sizing code_cell rendered">
<div class="input">

<div class="inner_cell">
    <div class="input_area">
<div class=" highlight hl-ipython3"><pre><span></span><span class="n">model2</span><span class="o">.</span><span class="n">key_padding_mask</span> <span class="o">=</span> <span class="kc">True</span>
<span class="n">model2</span><span class="o">.</span><span class="n">to</span><span class="p">(</span><span class="n">device</span><span class="p">)((</span><span class="n">xb</span><span class="p">,</span> <span class="n">key_padding_mask</span><span class="p">))</span><span class="o">.</span><span class="n">shape</span>
</pre></div>

    </div>
</div>
</div>

<div class="output_wrapper">
<div class="output">

<div class="output_area">



<div class="output_text output_subarea output_execute_result">
<pre>torch.Size([8, 2])</pre>
</div>

</div>

</div>
</div>

</div>
    {% endraw %}

    {% raw %}
    
<div class="cell border-box-sizing code_cell rendered">
<div class="input">

<div class="inner_cell">
    <div class="input_area">
<div class=" highlight hl-ipython3"><pre><span></span><span class="n">model</span><span class="o">.</span><span class="n">head</span>
</pre></div>

    </div>
</div>
</div>

<div class="output_wrapper">
<div class="output">

<div class="output_area">



<div class="output_text output_subarea output_execute_result">
<pre>Sequential(
  (0): GELU()
  (1): Flatten(full=False)
  (2): LinBnDrop(
    (0): Dropout(p=0.1, inplace=False)
    (1): Linear(in_features=32768, out_features=2, bias=True)
  )
)</pre>
</div>

</div>

</div>
</div>

</div>
    {% endraw %}

    {% raw %}
    
<div class="cell border-box-sizing code_cell rendered">
<div class="input">

<div class="inner_cell">
    <div class="input_area">
<div class=" highlight hl-ipython3"><pre><span></span><span class="n">model</span> <span class="o">=</span> <span class="n">TSTPlus</span><span class="p">(</span><span class="n">c_in</span><span class="p">,</span> <span class="n">c_out</span><span class="p">,</span> <span class="n">seq_len</span><span class="p">,</span> <span class="n">pre_norm</span><span class="o">=</span><span class="kc">True</span><span class="p">)</span>
<span class="n">test_eq</span><span class="p">(</span><span class="n">model</span><span class="o">.</span><span class="n">to</span><span class="p">(</span><span class="n">xb</span><span class="o">.</span><span class="n">device</span><span class="p">)(</span><span class="n">xb</span><span class="p">)</span><span class="o">.</span><span class="n">shape</span><span class="p">,</span> <span class="p">[</span><span class="n">bs</span><span class="p">,</span> <span class="n">c_out</span><span class="p">])</span>
</pre></div>

    </div>
</div>
</div>

</div>
    {% endraw %}

    {% raw %}
    
<div class="cell border-box-sizing code_cell rendered">
<div class="input">

<div class="inner_cell">
    <div class="input_area">
<div class=" highlight hl-ipython3"><pre><span></span><span class="n">bs</span> <span class="o">=</span> <span class="mi">8</span>
<span class="n">c_in</span> <span class="o">=</span> <span class="mi">9</span>  <span class="c1"># aka channels, features, variables, dimensions</span>
<span class="n">c_out</span> <span class="o">=</span> <span class="mi">2</span>
<span class="n">seq_len</span> <span class="o">=</span> <span class="mi">5000</span>

<span class="n">xb</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">randn</span><span class="p">(</span><span class="n">bs</span><span class="p">,</span> <span class="n">c_in</span><span class="p">,</span> <span class="n">seq_len</span><span class="p">)</span>

<span class="c1"># standardize by channel by_var based on the training set</span>
<span class="n">xb</span> <span class="o">=</span> <span class="p">(</span><span class="n">xb</span> <span class="o">-</span> <span class="n">xb</span><span class="o">.</span><span class="n">mean</span><span class="p">((</span><span class="mi">0</span><span class="p">,</span> <span class="mi">2</span><span class="p">),</span> <span class="n">keepdim</span><span class="o">=</span><span class="kc">True</span><span class="p">))</span> <span class="o">/</span> <span class="n">xb</span><span class="o">.</span><span class="n">std</span><span class="p">((</span><span class="mi">0</span><span class="p">,</span> <span class="mi">2</span><span class="p">),</span> <span class="n">keepdim</span><span class="o">=</span><span class="kc">True</span><span class="p">)</span>

<span class="n">model</span> <span class="o">=</span> <span class="n">TSTPlus</span><span class="p">(</span><span class="n">c_in</span><span class="p">,</span> <span class="n">c_out</span><span class="p">,</span> <span class="n">seq_len</span><span class="p">,</span> <span class="n">res_attention</span><span class="o">=</span><span class="kc">True</span><span class="p">)</span>
<span class="n">test_eq</span><span class="p">(</span><span class="n">model</span><span class="o">.</span><span class="n">to</span><span class="p">(</span><span class="n">xb</span><span class="o">.</span><span class="n">device</span><span class="p">)(</span><span class="n">xb</span><span class="p">)</span><span class="o">.</span><span class="n">shape</span><span class="p">,</span> <span class="p">[</span><span class="n">bs</span><span class="p">,</span> <span class="n">c_out</span><span class="p">])</span>
<span class="nb">print</span><span class="p">(</span><span class="sa">f</span><span class="s1">&#39;model parameters: </span><span class="si">{</span><span class="n">count_parameters</span><span class="p">(</span><span class="n">model</span><span class="p">)</span><span class="si">}</span><span class="s1">&#39;</span><span class="p">)</span>
</pre></div>

    </div>
</div>
</div>

<div class="output_wrapper">
<div class="output">

<div class="output_area">

<div class="output_subarea output_stream output_stdout output_text">
<pre>model parameters: 605698
</pre>
</div>
</div>

</div>
</div>

</div>
    {% endraw %}

    {% raw %}
    
<div class="cell border-box-sizing code_cell rendered">
<div class="input">

<div class="inner_cell">
    <div class="input_area">
<div class=" highlight hl-ipython3"><pre><span></span><span class="n">custom_head</span> <span class="o">=</span> <span class="n">partial</span><span class="p">(</span><span class="n">create_pool_head</span><span class="p">,</span> <span class="n">concat_pool</span><span class="o">=</span><span class="kc">True</span><span class="p">)</span>
<span class="n">model</span> <span class="o">=</span> <span class="n">TSTPlus</span><span class="p">(</span><span class="n">c_in</span><span class="p">,</span> <span class="n">c_out</span><span class="p">,</span> <span class="n">seq_len</span><span class="p">,</span> <span class="n">max_seq_len</span><span class="o">=</span><span class="n">max_seq_len</span><span class="p">,</span> <span class="n">d_model</span><span class="o">=</span><span class="n">d_model</span><span class="p">,</span> <span class="n">n_heads</span><span class="o">=</span><span class="n">n_heads</span><span class="p">,</span>
            <span class="n">d_k</span><span class="o">=</span><span class="n">d_k</span><span class="p">,</span> <span class="n">d_v</span><span class="o">=</span><span class="n">d_v</span><span class="p">,</span> <span class="n">d_ff</span><span class="o">=</span><span class="n">d_ff</span><span class="p">,</span> <span class="n">dropout</span><span class="o">=</span><span class="n">dropout</span><span class="p">,</span> <span class="n">activation</span><span class="o">=</span><span class="n">activation</span><span class="p">,</span> <span class="n">n_layers</span><span class="o">=</span><span class="n">n_layers</span><span class="p">,</span>
            <span class="n">fc_dropout</span><span class="o">=</span><span class="n">fc_dropout</span><span class="p">,</span> <span class="n">pe</span><span class="o">=</span><span class="n">pe</span><span class="p">,</span> <span class="n">learn_pe</span><span class="o">=</span><span class="n">learn_pe</span><span class="p">,</span> <span class="n">flatten</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span> <span class="n">custom_head</span><span class="o">=</span><span class="n">custom_head</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">)</span>
<span class="n">test_eq</span><span class="p">(</span><span class="n">model</span><span class="o">.</span><span class="n">to</span><span class="p">(</span><span class="n">xb</span><span class="o">.</span><span class="n">device</span><span class="p">)(</span><span class="n">xb</span><span class="p">)</span><span class="o">.</span><span class="n">shape</span><span class="p">,</span> <span class="p">[</span><span class="n">bs</span><span class="p">,</span> <span class="n">c_out</span><span class="p">])</span>
<span class="nb">print</span><span class="p">(</span><span class="sa">f</span><span class="s1">&#39;model parameters: </span><span class="si">{</span><span class="n">count_parameters</span><span class="p">(</span><span class="n">model</span><span class="p">)</span><span class="si">}</span><span class="s1">&#39;</span><span class="p">)</span>
</pre></div>

    </div>
</div>
</div>

<div class="output_wrapper">
<div class="output">

<div class="output_area">

<div class="output_subarea output_stream output_stdout output_text">
<pre>model parameters: 421122
</pre>
</div>
</div>

</div>
</div>

</div>
    {% endraw %}

    {% raw %}
    
<div class="cell border-box-sizing code_cell rendered">
<div class="input">

<div class="inner_cell">
    <div class="input_area">
<div class=" highlight hl-ipython3"><pre><span></span><span class="n">custom_head</span> <span class="o">=</span> <span class="n">partial</span><span class="p">(</span><span class="n">create_pool_plus_head</span><span class="p">,</span> <span class="n">concat_pool</span><span class="o">=</span><span class="kc">True</span><span class="p">)</span>
<span class="n">model</span> <span class="o">=</span> <span class="n">TSTPlus</span><span class="p">(</span><span class="n">c_in</span><span class="p">,</span> <span class="n">c_out</span><span class="p">,</span> <span class="n">seq_len</span><span class="p">,</span> <span class="n">max_seq_len</span><span class="o">=</span><span class="n">max_seq_len</span><span class="p">,</span> <span class="n">d_model</span><span class="o">=</span><span class="n">d_model</span><span class="p">,</span> <span class="n">n_heads</span><span class="o">=</span><span class="n">n_heads</span><span class="p">,</span>
            <span class="n">d_k</span><span class="o">=</span><span class="n">d_k</span><span class="p">,</span> <span class="n">d_v</span><span class="o">=</span><span class="n">d_v</span><span class="p">,</span> <span class="n">d_ff</span><span class="o">=</span><span class="n">d_ff</span><span class="p">,</span> <span class="n">dropout</span><span class="o">=</span><span class="n">dropout</span><span class="p">,</span> <span class="n">activation</span><span class="o">=</span><span class="n">activation</span><span class="p">,</span> <span class="n">n_layers</span><span class="o">=</span><span class="n">n_layers</span><span class="p">,</span>
            <span class="n">fc_dropout</span><span class="o">=</span><span class="n">fc_dropout</span><span class="p">,</span> <span class="n">pe</span><span class="o">=</span><span class="n">pe</span><span class="p">,</span> <span class="n">learn_pe</span><span class="o">=</span><span class="n">learn_pe</span><span class="p">,</span> <span class="n">flatten</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span> <span class="n">custom_head</span><span class="o">=</span><span class="n">custom_head</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">)</span>
<span class="n">test_eq</span><span class="p">(</span><span class="n">model</span><span class="o">.</span><span class="n">to</span><span class="p">(</span><span class="n">xb</span><span class="o">.</span><span class="n">device</span><span class="p">)(</span><span class="n">xb</span><span class="p">)</span><span class="o">.</span><span class="n">shape</span><span class="p">,</span> <span class="p">[</span><span class="n">bs</span><span class="p">,</span> <span class="n">c_out</span><span class="p">])</span>
<span class="nb">print</span><span class="p">(</span><span class="sa">f</span><span class="s1">&#39;model parameters: </span><span class="si">{</span><span class="n">count_parameters</span><span class="p">(</span><span class="n">model</span><span class="p">)</span><span class="si">}</span><span class="s1">&#39;</span><span class="p">)</span>
</pre></div>

    </div>
</div>
</div>

<div class="output_wrapper">
<div class="output">

<div class="output_area">

<div class="output_subarea output_stream output_stdout output_text">
<pre>model parameters: 554240
</pre>
</div>
</div>

</div>
</div>

</div>
    {% endraw %}

    {% raw %}
    
<div class="cell border-box-sizing code_cell rendered">
<div class="input">

<div class="inner_cell">
    <div class="input_area">
<div class=" highlight hl-ipython3"><pre><span></span><span class="n">bs</span> <span class="o">=</span> <span class="mi">8</span>
<span class="n">c_in</span> <span class="o">=</span> <span class="mi">9</span>  <span class="c1"># aka channels, features, variables, dimensions</span>
<span class="n">c_out</span> <span class="o">=</span> <span class="mi">2</span>
<span class="n">seq_len</span> <span class="o">=</span> <span class="mi">60</span>

<span class="n">xb</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">randn</span><span class="p">(</span><span class="n">bs</span><span class="p">,</span> <span class="n">c_in</span><span class="p">,</span> <span class="n">seq_len</span><span class="p">)</span>

<span class="c1"># standardize by channel by_var based on the training set</span>
<span class="n">xb</span> <span class="o">=</span> <span class="p">(</span><span class="n">xb</span> <span class="o">-</span> <span class="n">xb</span><span class="o">.</span><span class="n">mean</span><span class="p">((</span><span class="mi">0</span><span class="p">,</span> <span class="mi">2</span><span class="p">),</span> <span class="n">keepdim</span><span class="o">=</span><span class="kc">True</span><span class="p">))</span> <span class="o">/</span> <span class="n">xb</span><span class="o">.</span><span class="n">std</span><span class="p">((</span><span class="mi">0</span><span class="p">,</span> <span class="mi">2</span><span class="p">),</span> <span class="n">keepdim</span><span class="o">=</span><span class="kc">True</span><span class="p">)</span>

<span class="c1"># Settings</span>
<span class="n">max_seq_len</span> <span class="o">=</span> <span class="mi">120</span>
<span class="n">d_model</span> <span class="o">=</span> <span class="mi">128</span>
<span class="n">n_heads</span> <span class="o">=</span> <span class="mi">16</span>
<span class="n">d_k</span> <span class="o">=</span> <span class="n">d_v</span> <span class="o">=</span> <span class="kc">None</span> <span class="c1"># if None --&gt; d_model // n_heads</span>
<span class="n">d_ff</span> <span class="o">=</span> <span class="mi">256</span>
<span class="n">dropout</span> <span class="o">=</span> <span class="mf">0.1</span>
<span class="n">act</span> <span class="o">=</span> <span class="s2">&quot;gelu&quot;</span>
<span class="n">n_layers</span> <span class="o">=</span> <span class="mi">3</span>
<span class="n">fc_dropout</span> <span class="o">=</span> <span class="mf">0.1</span>
<span class="n">pe</span><span class="o">=</span><span class="s1">&#39;zeros&#39;</span>
<span class="n">learn_pe</span><span class="o">=</span><span class="kc">True</span>
<span class="n">kwargs</span> <span class="o">=</span> <span class="p">{}</span>
<span class="c1"># kwargs = dict(kernel_size=5, padding=2)</span>

<span class="n">model</span> <span class="o">=</span> <span class="n">TSTPlus</span><span class="p">(</span><span class="n">c_in</span><span class="p">,</span> <span class="n">c_out</span><span class="p">,</span> <span class="n">seq_len</span><span class="p">,</span> <span class="n">max_seq_len</span><span class="o">=</span><span class="n">max_seq_len</span><span class="p">,</span> <span class="n">d_model</span><span class="o">=</span><span class="n">d_model</span><span class="p">,</span> <span class="n">n_heads</span><span class="o">=</span><span class="n">n_heads</span><span class="p">,</span>
            <span class="n">d_k</span><span class="o">=</span><span class="n">d_k</span><span class="p">,</span> <span class="n">d_v</span><span class="o">=</span><span class="n">d_v</span><span class="p">,</span> <span class="n">d_ff</span><span class="o">=</span><span class="n">d_ff</span><span class="p">,</span> <span class="n">dropout</span><span class="o">=</span><span class="n">dropout</span><span class="p">,</span> <span class="n">act</span><span class="o">=</span><span class="n">act</span><span class="p">,</span> <span class="n">n_layers</span><span class="o">=</span><span class="n">n_layers</span><span class="p">,</span>
            <span class="n">fc_dropout</span><span class="o">=</span><span class="n">fc_dropout</span><span class="p">,</span> <span class="n">pe</span><span class="o">=</span><span class="n">pe</span><span class="p">,</span> <span class="n">learn_pe</span><span class="o">=</span><span class="n">learn_pe</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">)</span>
<span class="n">test_eq</span><span class="p">(</span><span class="n">model</span><span class="o">.</span><span class="n">to</span><span class="p">(</span><span class="n">xb</span><span class="o">.</span><span class="n">device</span><span class="p">)(</span><span class="n">xb</span><span class="p">)</span><span class="o">.</span><span class="n">shape</span><span class="p">,</span> <span class="p">[</span><span class="n">bs</span><span class="p">,</span> <span class="n">c_out</span><span class="p">])</span>
<span class="nb">print</span><span class="p">(</span><span class="sa">f</span><span class="s1">&#39;model parameters: </span><span class="si">{</span><span class="n">count_parameters</span><span class="p">(</span><span class="n">model</span><span class="p">)</span><span class="si">}</span><span class="s1">&#39;</span><span class="p">)</span>
<span class="n">body</span><span class="p">,</span> <span class="n">head</span> <span class="o">=</span> <span class="n">model</span><span class="p">[</span><span class="mi">0</span><span class="p">],</span> <span class="n">model</span><span class="p">[</span><span class="mi">1</span><span class="p">]</span>
<span class="n">test_eq</span><span class="p">(</span><span class="n">body</span><span class="o">.</span><span class="n">to</span><span class="p">(</span><span class="n">xb</span><span class="o">.</span><span class="n">device</span><span class="p">)(</span><span class="n">xb</span><span class="p">)</span><span class="o">.</span><span class="n">ndim</span><span class="p">,</span> <span class="mi">3</span><span class="p">)</span>
<span class="n">test_eq</span><span class="p">(</span><span class="n">head</span><span class="o">.</span><span class="n">to</span><span class="p">(</span><span class="n">xb</span><span class="o">.</span><span class="n">device</span><span class="p">)(</span><span class="n">body</span><span class="o">.</span><span class="n">to</span><span class="p">(</span><span class="n">xb</span><span class="o">.</span><span class="n">device</span><span class="p">)(</span><span class="n">xb</span><span class="p">))</span><span class="o">.</span><span class="n">ndim</span><span class="p">,</span> <span class="mi">2</span><span class="p">)</span>
<span class="n">head</span>
</pre></div>

    </div>
</div>
</div>

<div class="output_wrapper">
<div class="output">

<div class="output_area">

<div class="output_subarea output_stream output_stdout output_text">
<pre>model parameters: 421762
</pre>
</div>
</div>

<div class="output_area">



<div class="output_text output_subarea output_execute_result">
<pre>Sequential(
  (0): GELU()
  (1): Flatten(full=False)
  (2): LinBnDrop(
    (0): Dropout(p=0.1, inplace=False)
    (1): Linear(in_features=7680, out_features=2, bias=True)
  )
)</pre>
</div>

</div>

</div>
</div>

</div>
    {% endraw %}

    {% raw %}
    
<div class="cell border-box-sizing code_cell rendered">
<div class="input">

<div class="inner_cell">
    <div class="input_area">
<div class=" highlight hl-ipython3"><pre><span></span><span class="n">model</span><span class="o">.</span><span class="n">show_pe</span><span class="p">()</span>
</pre></div>

    </div>
</div>
</div>

<div class="output_wrapper">
<div class="output">

<div class="output_area">



<div class="output_png output_subarea ">
<img src="
"
>
</div>

</div>

<div class="output_area">



<div class="output_png output_subarea ">
<img src="
"
>
</div>

</div>

</div>
</div>

</div>
    {% endraw %}

    {% raw %}
    
<div class="cell border-box-sizing code_cell rendered">
<div class="input">

<div class="inner_cell">
    <div class="input_area">
<div class=" highlight hl-ipython3"><pre><span></span><span class="n">model</span> <span class="o">=</span> <span class="n">TSTPlus</span><span class="p">(</span><span class="mi">3</span><span class="p">,</span> <span class="mi">2</span><span class="p">,</span> <span class="mi">10</span><span class="p">)</span>
<span class="n">xb</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">randn</span><span class="p">(</span><span class="mi">4</span><span class="p">,</span> <span class="mi">3</span><span class="p">,</span> <span class="mi">10</span><span class="p">)</span>
<span class="n">yb</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">randint</span><span class="p">(</span><span class="mi">0</span><span class="p">,</span> <span class="mi">2</span><span class="p">,</span> <span class="p">(</span><span class="mi">4</span><span class="p">,))</span>
<span class="n">test_eq</span><span class="p">(</span><span class="n">model</span><span class="o">.</span><span class="n">backbone</span><span class="o">.</span><span class="n">_key_padding_mask</span><span class="p">(</span><span class="n">xb</span><span class="p">)[</span><span class="mi">1</span><span class="p">],</span> <span class="kc">None</span><span class="p">)</span>
<span class="n">random_idxs</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">random</span><span class="o">.</span><span class="n">choice</span><span class="p">(</span><span class="nb">len</span><span class="p">(</span><span class="n">xb</span><span class="p">),</span> <span class="mi">2</span><span class="p">,</span> <span class="kc">False</span><span class="p">)</span>
<span class="n">xb</span><span class="p">[</span><span class="n">random_idxs</span><span class="p">,</span> <span class="p">:,</span> <span class="o">-</span><span class="mi">5</span><span class="p">:]</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">nan</span>
<span class="n">xb</span><span class="p">[</span><span class="n">random_idxs</span><span class="p">,</span> <span class="mi">0</span><span class="p">,</span> <span class="mi">1</span><span class="p">]</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">nan</span>
<span class="n">test_eq</span><span class="p">(</span><span class="n">model</span><span class="o">.</span><span class="n">backbone</span><span class="o">.</span><span class="n">_key_padding_mask</span><span class="p">(</span><span class="n">xb</span><span class="o">.</span><span class="n">clone</span><span class="p">())[</span><span class="mi">1</span><span class="p">]</span><span class="o">.</span><span class="n">data</span><span class="p">,</span> <span class="p">(</span><span class="n">torch</span><span class="o">.</span><span class="n">isnan</span><span class="p">(</span><span class="n">xb</span><span class="p">)</span><span class="o">.</span><span class="n">float</span><span class="p">()</span><span class="o">.</span><span class="n">mean</span><span class="p">(</span><span class="mi">1</span><span class="p">)</span><span class="o">==</span><span class="mi">1</span><span class="p">)</span><span class="o">.</span><span class="n">bool</span><span class="p">())</span>
<span class="n">test_eq</span><span class="p">(</span><span class="n">model</span><span class="o">.</span><span class="n">backbone</span><span class="o">.</span><span class="n">_key_padding_mask</span><span class="p">(</span><span class="n">xb</span><span class="o">.</span><span class="n">clone</span><span class="p">())[</span><span class="mi">1</span><span class="p">]</span><span class="o">.</span><span class="n">data</span><span class="o">.</span><span class="n">shape</span><span class="p">,</span> <span class="p">(</span><span class="mi">4</span><span class="p">,</span><span class="mi">10</span><span class="p">))</span>
<span class="nb">print</span><span class="p">(</span><span class="n">torch</span><span class="o">.</span><span class="n">isnan</span><span class="p">(</span><span class="n">xb</span><span class="p">)</span><span class="o">.</span><span class="n">sum</span><span class="p">())</span>
<span class="n">pred</span> <span class="o">=</span> <span class="n">model</span><span class="o">.</span><span class="n">to</span><span class="p">(</span><span class="n">xb</span><span class="o">.</span><span class="n">device</span><span class="p">)(</span><span class="n">xb</span><span class="o">.</span><span class="n">clone</span><span class="p">())</span>
<span class="n">loss</span> <span class="o">=</span> <span class="n">CrossEntropyLossFlat</span><span class="p">()(</span><span class="n">pred</span><span class="p">,</span> <span class="n">yb</span><span class="p">)</span>
<span class="n">loss</span><span class="o">.</span><span class="n">backward</span><span class="p">()</span>
<span class="n">model</span><span class="o">.</span><span class="n">to</span><span class="p">(</span><span class="n">xb</span><span class="o">.</span><span class="n">device</span><span class="p">)</span><span class="o">.</span><span class="n">backbone</span><span class="o">.</span><span class="n">_key_padding_mask</span><span class="p">(</span><span class="n">xb</span><span class="p">)[</span><span class="mi">1</span><span class="p">]</span><span class="o">.</span><span class="n">data</span><span class="o">.</span><span class="n">shape</span>
</pre></div>

    </div>
</div>
</div>

<div class="output_wrapper">
<div class="output">

<div class="output_area">

<div class="output_subarea output_stream output_stdout output_text">
<pre>tensor(32)
</pre>
</div>
</div>

<div class="output_area">



<div class="output_text output_subarea output_execute_result">
<pre>torch.Size([4, 10])</pre>
</div>

</div>

</div>
</div>

</div>
    {% endraw %}

    {% raw %}
    
<div class="cell border-box-sizing code_cell rendered">
<div class="input">

<div class="inner_cell">
    <div class="input_area">
<div class=" highlight hl-ipython3"><pre><span></span><span class="n">bs</span> <span class="o">=</span> <span class="mi">4</span>
<span class="n">c_in</span> <span class="o">=</span> <span class="mi">3</span>
<span class="n">seq_len</span> <span class="o">=</span> <span class="mi">10</span>
<span class="n">c_out</span> <span class="o">=</span> <span class="mi">2</span>
<span class="n">xb</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">randn</span><span class="p">(</span><span class="n">bs</span><span class="p">,</span> <span class="n">c_in</span><span class="p">,</span> <span class="n">seq_len</span><span class="p">)</span>
<span class="n">xb</span><span class="p">[:,</span> <span class="o">-</span><span class="mi">1</span><span class="p">]</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">randint</span><span class="p">(</span><span class="mi">0</span><span class="p">,</span> <span class="mi">2</span><span class="p">,</span> <span class="p">(</span><span class="n">bs</span><span class="p">,</span> <span class="n">seq_len</span><span class="p">))</span><span class="o">.</span><span class="n">sort</span><span class="p">()[</span><span class="mi">0</span><span class="p">]</span>
<span class="n">model</span> <span class="o">=</span> <span class="n">TSTPlus</span><span class="p">(</span><span class="n">c_in</span><span class="p">,</span> <span class="n">c_out</span><span class="p">,</span> <span class="n">seq_len</span><span class="p">)</span><span class="o">.</span><span class="n">to</span><span class="p">(</span><span class="n">xb</span><span class="o">.</span><span class="n">device</span><span class="p">)</span>
<span class="n">test_eq</span><span class="p">(</span><span class="n">model</span><span class="o">.</span><span class="n">backbone</span><span class="o">.</span><span class="n">_key_padding_mask</span><span class="p">(</span><span class="n">xb</span><span class="p">)[</span><span class="mi">1</span><span class="p">],</span> <span class="kc">None</span><span class="p">)</span>
<span class="n">model</span> <span class="o">=</span> <span class="n">TSTPlus</span><span class="p">(</span><span class="n">c_in</span><span class="p">,</span> <span class="n">c_out</span><span class="p">,</span> <span class="n">seq_len</span><span class="p">,</span> <span class="n">padding_var</span><span class="o">=-</span><span class="mi">1</span><span class="p">)</span><span class="o">.</span><span class="n">to</span><span class="p">(</span><span class="n">xb</span><span class="o">.</span><span class="n">device</span><span class="p">)</span>
<span class="n">test_eq</span><span class="p">(</span><span class="n">model</span><span class="o">.</span><span class="n">backbone</span><span class="o">.</span><span class="n">_key_padding_mask</span><span class="p">(</span><span class="n">xb</span><span class="p">)[</span><span class="mi">1</span><span class="p">],</span> <span class="p">(</span><span class="n">xb</span><span class="p">[:,</span> <span class="o">-</span><span class="mi">1</span><span class="p">]</span><span class="o">==</span><span class="mi">1</span><span class="p">))</span>
<span class="n">model</span> <span class="o">=</span> <span class="n">TSTPlus</span><span class="p">(</span><span class="n">c_in</span><span class="p">,</span> <span class="n">c_out</span><span class="p">,</span> <span class="n">seq_len</span><span class="p">,</span> <span class="n">padding_var</span><span class="o">=</span><span class="mi">2</span><span class="p">)</span><span class="o">.</span><span class="n">to</span><span class="p">(</span><span class="n">xb</span><span class="o">.</span><span class="n">device</span><span class="p">)</span>
<span class="n">test_eq</span><span class="p">(</span><span class="n">model</span><span class="o">.</span><span class="n">backbone</span><span class="o">.</span><span class="n">_key_padding_mask</span><span class="p">(</span><span class="n">xb</span><span class="p">)[</span><span class="mi">1</span><span class="p">],</span> <span class="p">(</span><span class="n">xb</span><span class="p">[:,</span> <span class="o">-</span><span class="mi">1</span><span class="p">]</span><span class="o">==</span><span class="mi">1</span><span class="p">))</span>
<span class="n">test_eq</span><span class="p">(</span><span class="n">model</span><span class="p">(</span><span class="n">xb</span><span class="p">)</span><span class="o">.</span><span class="n">shape</span><span class="p">,</span> <span class="p">(</span><span class="n">bs</span><span class="p">,</span> <span class="n">c_out</span><span class="p">))</span>
</pre></div>

    </div>
</div>
</div>

</div>
    {% endraw %}

    {% raw %}
    
<div class="cell border-box-sizing code_cell rendered">

<div class="output_wrapper">
<div class="output">

<div class="output_area">


<div class="output_markdown rendered_html output_subarea ">
<h2 id="MultiTSTPlus" class="doc_header"><code>class</code> <code>MultiTSTPlus</code><a href="https://github.com/timeseriesAI/tsai/tree/main/tsai/models/TSTPlus.py#L350" class="source_link" style="float:right">[source]</a></h2><blockquote><p><code>MultiTSTPlus</code>(<strong><code>feat_list</code></strong>, <strong><code>c_out</code></strong>, <strong><code>seq_len</code></strong>, <strong><code>max_seq_len</code></strong>:<code>Optional</code>[<code>int</code>]=<em><code>512</code></em>, <strong><code>custom_head</code></strong>=<em><code>None</code></em>, <strong><code>n_layers</code></strong>:<code>int</code>=<em><code>3</code></em>, <strong><code>d_model</code></strong>:<code>int</code>=<em><code>128</code></em>, <strong><code>n_heads</code></strong>:<code>int</code>=<em><code>16</code></em>, <strong><code>d_k</code></strong>:<code>Optional</code>[<code>int</code>]=<em><code>None</code></em>, <strong><code>d_v</code></strong>:<code>Optional</code>[<code>int</code>]=<em><code>None</code></em>, <strong><code>d_ff</code></strong>:<code>int</code>=<em><code>256</code></em>, <strong><code>norm</code></strong>:<code>str</code>=<em><code>'BatchNorm'</code></em>, <strong><code>attn_dropout</code></strong>:<code>float</code>=<em><code>0.0</code></em>, <strong><code>dropout</code></strong>:<code>float</code>=<em><code>0.0</code></em>, <strong><code>act</code></strong>:<code>str</code>=<em><code>'gelu'</code></em>, <strong><code>key_padding_mask</code></strong>:<code>bool</code>=<em><code>'auto'</code></em>, <strong><code>padding_var</code></strong>:<code>Optional</code>[<code>int</code>]=<em><code>None</code></em>, <strong><code>attn_mask</code></strong>:<code>Optional</code>[<code>Tensor</code>]=<em><code>None</code></em>, <strong><code>res_attention</code></strong>:<code>bool</code>=<em><code>True</code></em>, <strong><code>pre_norm</code></strong>:<code>bool</code>=<em><code>False</code></em>, <strong><code>store_attn</code></strong>:<code>bool</code>=<em><code>False</code></em>, <strong><code>pe</code></strong>:<code>str</code>=<em><code>'zeros'</code></em>, <strong><code>learn_pe</code></strong>:<code>bool</code>=<em><code>True</code></em>, <strong><code>flatten</code></strong>:<code>bool</code>=<em><code>True</code></em>, <strong><code>fc_dropout</code></strong>:<code>float</code>=<em><code>0.0</code></em>, <strong><code>concat_pool</code></strong>:<code>bool</code>=<em><code>False</code></em>, <strong><code>bn</code></strong>:<code>bool</code>=<em><code>False</code></em>, <strong><code>y_range</code></strong>:<code>Optional</code>[<code>tuple</code>]=<em><code>None</code></em>, <strong><code>verbose</code></strong>:<code>bool</code>=<em><code>False</code></em>) :: <a href="/models.TabFusionTransformer.html#Sequential"><code>Sequential</code></a></p>
</blockquote>
<p>A sequential container.
Modules will be added to it in the order they are passed in the
constructor. Alternatively, an <code>OrderedDict</code> of modules can be
passed in. The <code>forward()</code> method of <code>Sequential</code> accepts any
input and forwards it to the first module it contains. It then
"chains" outputs to inputs sequentially for each subsequent module,
finally returning the output of the last module.</p>
<p>The value a <code>Sequential</code> provides over manually calling a sequence
of modules is that it allows treating the whole container as a
single module, such that performing a transformation on the
<code>Sequential</code> applies to each of the modules it stores (which are
each a registered submodule of the <code>Sequential</code>).</p>
<p>What's the difference between a <code>Sequential</code> and a
:class:<code>torch.nn.ModuleList</code>? A <code>ModuleList</code> is exactly what it
sounds like--a list for storing <code>Module</code> s! On the other hand,
the layers in a <code>Sequential</code> are connected in a cascading way.</p>
<p>Example::</p>

<pre><code># Using Sequential to create a small model. When `model` is run,
# input will first be passed to `Conv2d(1,20,5)`. The output of
# `Conv2d(1,20,5)` will be used as the input to the first
# `ReLU`; the output of the first `ReLU` will become the input
# for `Conv2d(20,64,5)`. Finally, the output of
# `Conv2d(20,64,5)` will be used as input to the second `ReLU`
model = nn.Sequential(
          nn.Conv2d(1,20,5),
          nn.ReLU(),
          nn.Conv2d(20,64,5),
          nn.ReLU()
        )

# Using Sequential with OrderedDict. This is functionally the
# same as the above code
model = nn.Sequential(OrderedDict([
          ('conv1', nn.Conv2d(1,20,5)),
          ('relu1', nn.ReLU()),
          ('conv2', nn.Conv2d(20,64,5)),
          ('relu2', nn.ReLU())
        ]))</code></pre>

</div>

</div>

</div>
</div>

</div>
    {% endraw %}

    {% raw %}
    
<div class="cell border-box-sizing code_cell rendered">

</div>
    {% endraw %}

    {% raw %}
    
<div class="cell border-box-sizing code_cell rendered">
<div class="input">

<div class="inner_cell">
    <div class="input_area">
<div class=" highlight hl-ipython3"><pre><span></span><span class="n">bs</span> <span class="o">=</span> <span class="mi">8</span>
<span class="n">c_in</span> <span class="o">=</span> <span class="mi">7</span>  <span class="c1"># aka channels, features, variables, dimensions</span>
<span class="n">c_out</span> <span class="o">=</span> <span class="mi">2</span>
<span class="n">seq_len</span> <span class="o">=</span> <span class="mi">10</span>
<span class="n">xb2</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">randn</span><span class="p">(</span><span class="n">bs</span><span class="p">,</span> <span class="n">c_in</span><span class="p">,</span> <span class="n">seq_len</span><span class="p">)</span>
<span class="n">model1</span> <span class="o">=</span> <span class="n">MultiTSTPlus</span><span class="p">([</span><span class="mi">2</span><span class="p">,</span> <span class="mi">5</span><span class="p">],</span> <span class="n">c_out</span><span class="p">,</span> <span class="n">seq_len</span><span class="p">)</span>
<span class="n">model2</span> <span class="o">=</span> <span class="n">MultiTSTPlus</span><span class="p">(</span><span class="mi">7</span><span class="p">,</span> <span class="n">c_out</span><span class="p">,</span> <span class="n">seq_len</span><span class="p">)</span>
<span class="n">test_eq</span><span class="p">(</span><span class="n">model1</span><span class="o">.</span><span class="n">to</span><span class="p">(</span><span class="n">xb2</span><span class="o">.</span><span class="n">device</span><span class="p">)(</span><span class="n">xb2</span><span class="p">)</span><span class="o">.</span><span class="n">shape</span><span class="p">,</span> <span class="p">(</span><span class="n">bs</span><span class="p">,</span> <span class="n">c_out</span><span class="p">))</span>
<span class="n">test_eq</span><span class="p">(</span><span class="n">model1</span><span class="o">.</span><span class="n">to</span><span class="p">(</span><span class="n">xb2</span><span class="o">.</span><span class="n">device</span><span class="p">)(</span><span class="n">xb2</span><span class="p">)</span><span class="o">.</span><span class="n">shape</span><span class="p">,</span> <span class="n">model2</span><span class="o">.</span><span class="n">to</span><span class="p">(</span><span class="n">xb2</span><span class="o">.</span><span class="n">device</span><span class="p">)(</span><span class="n">xb2</span><span class="p">)</span><span class="o">.</span><span class="n">shape</span><span class="p">)</span>
<span class="n">test_eq</span><span class="p">(</span><span class="n">count_parameters</span><span class="p">(</span><span class="n">model1</span><span class="p">)</span> <span class="o">&gt;</span> <span class="n">count_parameters</span><span class="p">(</span><span class="n">model2</span><span class="p">),</span> <span class="kc">True</span><span class="p">)</span>
</pre></div>

    </div>
</div>
</div>

</div>
    {% endraw %}

    {% raw %}
    
<div class="cell border-box-sizing code_cell rendered">
<div class="input">

<div class="inner_cell">
    <div class="input_area">
<div class=" highlight hl-ipython3"><pre><span></span><span class="n">bs</span> <span class="o">=</span> <span class="mi">8</span>
<span class="n">c_in</span> <span class="o">=</span> <span class="mi">7</span>  <span class="c1"># aka channels, features, variables, dimensions</span>
<span class="n">c_out</span> <span class="o">=</span> <span class="mi">2</span>
<span class="n">seq_len</span> <span class="o">=</span> <span class="mi">10</span>
<span class="n">xb2</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">randn</span><span class="p">(</span><span class="n">bs</span><span class="p">,</span> <span class="n">c_in</span><span class="p">,</span> <span class="n">seq_len</span><span class="p">)</span>
<span class="n">model1</span> <span class="o">=</span> <span class="n">MultiTSTPlus</span><span class="p">([</span><span class="mi">2</span><span class="p">,</span> <span class="mi">5</span><span class="p">],</span> <span class="n">c_out</span><span class="p">,</span> <span class="n">seq_len</span><span class="p">,</span> <span class="p">)</span>
<span class="n">model2</span> <span class="o">=</span> <span class="n">MultiTSTPlus</span><span class="p">([[</span><span class="mi">0</span><span class="p">,</span><span class="mi">2</span><span class="p">,</span><span class="mi">5</span><span class="p">],</span> <span class="p">[</span><span class="mi">0</span><span class="p">,</span><span class="mi">1</span><span class="p">,</span><span class="mi">3</span><span class="p">,</span><span class="mi">4</span><span class="p">,</span><span class="mi">6</span><span class="p">]],</span> <span class="n">c_out</span><span class="p">,</span> <span class="n">seq_len</span><span class="p">)</span>
<span class="n">test_eq</span><span class="p">(</span><span class="n">model1</span><span class="o">.</span><span class="n">to</span><span class="p">(</span><span class="n">xb2</span><span class="o">.</span><span class="n">device</span><span class="p">)(</span><span class="n">xb2</span><span class="p">)</span><span class="o">.</span><span class="n">shape</span><span class="p">,</span> <span class="p">(</span><span class="n">bs</span><span class="p">,</span> <span class="n">c_out</span><span class="p">))</span>
<span class="n">test_eq</span><span class="p">(</span><span class="n">model1</span><span class="o">.</span><span class="n">to</span><span class="p">(</span><span class="n">xb2</span><span class="o">.</span><span class="n">device</span><span class="p">)(</span><span class="n">xb2</span><span class="p">)</span><span class="o">.</span><span class="n">shape</span><span class="p">,</span> <span class="n">model2</span><span class="o">.</span><span class="n">to</span><span class="p">(</span><span class="n">xb2</span><span class="o">.</span><span class="n">device</span><span class="p">)(</span><span class="n">xb2</span><span class="p">)</span><span class="o">.</span><span class="n">shape</span><span class="p">)</span>
</pre></div>

    </div>
</div>
</div>

</div>
    {% endraw %}

    {% raw %}
    
<div class="cell border-box-sizing code_cell rendered">
<div class="input">

<div class="inner_cell">
    <div class="input_area">
<div class=" highlight hl-ipython3"><pre><span></span><span class="n">model1</span> <span class="o">=</span> <span class="n">MultiTSTPlus</span><span class="p">([</span><span class="mi">2</span><span class="p">,</span> <span class="mi">5</span><span class="p">],</span> <span class="n">c_out</span><span class="p">,</span> <span class="n">seq_len</span><span class="p">,</span> <span class="n">y_range</span><span class="o">=</span><span class="p">(</span><span class="mf">0.5</span><span class="p">,</span> <span class="mf">5.5</span><span class="p">))</span>
<span class="n">body</span><span class="p">,</span> <span class="n">head</span> <span class="o">=</span> <span class="n">split_model</span><span class="p">(</span><span class="n">model1</span><span class="p">)</span>
<span class="n">test_eq</span><span class="p">(</span><span class="n">body</span><span class="o">.</span><span class="n">to</span><span class="p">(</span><span class="n">xb2</span><span class="o">.</span><span class="n">device</span><span class="p">)(</span><span class="n">xb2</span><span class="p">)</span><span class="o">.</span><span class="n">ndim</span><span class="p">,</span> <span class="mi">3</span><span class="p">)</span>
<span class="n">test_eq</span><span class="p">(</span><span class="n">head</span><span class="o">.</span><span class="n">to</span><span class="p">(</span><span class="n">xb2</span><span class="o">.</span><span class="n">device</span><span class="p">)(</span><span class="n">body</span><span class="o">.</span><span class="n">to</span><span class="p">(</span><span class="n">xb2</span><span class="o">.</span><span class="n">device</span><span class="p">)(</span><span class="n">xb2</span><span class="p">))</span><span class="o">.</span><span class="n">ndim</span><span class="p">,</span> <span class="mi">2</span><span class="p">)</span>
<span class="n">head</span>
</pre></div>

    </div>
</div>
</div>

<div class="output_wrapper">
<div class="output">

<div class="output_area">



<div class="output_text output_subarea output_execute_result">
<pre>Sequential(
  (0): Sequential(
    (0): GELU()
    (1): Flatten(full=False)
    (2): LinBnDrop(
      (0): Linear(in_features=2560, out_features=2, bias=True)
    )
  )
)</pre>
</div>

</div>

</div>
</div>

</div>
    {% endraw %}

    {% raw %}
    
<div class="cell border-box-sizing code_cell rendered">
<div class="input">

<div class="inner_cell">
    <div class="input_area">
<div class=" highlight hl-ipython3"><pre><span></span><span class="n">model</span> <span class="o">=</span> <span class="n">MultiTSTPlus</span><span class="p">([</span><span class="mi">2</span><span class="p">,</span> <span class="mi">5</span><span class="p">],</span> <span class="n">c_out</span><span class="p">,</span> <span class="n">seq_len</span><span class="p">,</span> <span class="n">pre_norm</span><span class="o">=</span><span class="kc">True</span><span class="p">)</span>
</pre></div>

    </div>
</div>
</div>

</div>
    {% endraw %}

    {% raw %}
    
<div class="cell border-box-sizing code_cell rendered">
<div class="input">

<div class="inner_cell">
    <div class="input_area">
<div class=" highlight hl-ipython3"><pre><span></span><span class="n">bs</span> <span class="o">=</span> <span class="mi">8</span>
<span class="n">n_vars</span> <span class="o">=</span> <span class="mi">3</span>
<span class="n">seq_len</span> <span class="o">=</span> <span class="mi">12</span>
<span class="n">c_out</span> <span class="o">=</span> <span class="mi">2</span>
<span class="n">xb</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">rand</span><span class="p">(</span><span class="n">bs</span><span class="p">,</span> <span class="n">n_vars</span><span class="p">,</span> <span class="n">seq_len</span><span class="p">)</span>
<span class="n">net</span> <span class="o">=</span> <span class="n">MultiTSTPlus</span><span class="p">(</span><span class="n">n_vars</span><span class="p">,</span> <span class="n">c_out</span><span class="p">,</span> <span class="n">seq_len</span><span class="p">)</span>
<span class="n">change_model_head</span><span class="p">(</span><span class="n">net</span><span class="p">,</span> <span class="n">create_pool_plus_head</span><span class="p">,</span> <span class="n">concat_pool</span><span class="o">=</span><span class="kc">False</span><span class="p">)</span>
<span class="nb">print</span><span class="p">(</span><span class="n">net</span><span class="o">.</span><span class="n">to</span><span class="p">(</span><span class="n">xb</span><span class="o">.</span><span class="n">device</span><span class="p">)(</span><span class="n">xb</span><span class="p">)</span><span class="o">.</span><span class="n">shape</span><span class="p">)</span>
<span class="n">net</span><span class="o">.</span><span class="n">head</span>
</pre></div>

    </div>
</div>
</div>

<div class="output_wrapper">
<div class="output">

<div class="output_area">

<div class="output_subarea output_stream output_stdout output_text">
<pre>torch.Size([8, 2])
</pre>
</div>
</div>

<div class="output_area">



<div class="output_text output_subarea output_execute_result">
<pre>Sequential(
  (0): AdaptiveAvgPool1d(output_size=1)
  (1): Flatten(full=False)
  (2): BatchNorm1d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
  (3): Linear(in_features=128, out_features=512, bias=False)
  (4): ReLU(inplace=True)
  (5): BatchNorm1d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
  (6): Linear(in_features=512, out_features=2, bias=False)
)</pre>
</div>

</div>

</div>
</div>

</div>
    {% endraw %}

    {% raw %}
    
<div class="cell border-box-sizing code_cell rendered">
<div class="input">

<div class="inner_cell">
    <div class="input_area">
<div class=" highlight hl-ipython3"><pre><span></span><span class="n">bs</span> <span class="o">=</span> <span class="mi">8</span>
<span class="n">n_vars</span> <span class="o">=</span> <span class="mi">3</span>
<span class="n">seq_len</span> <span class="o">=</span> <span class="mi">12</span>
<span class="n">c_out</span> <span class="o">=</span> <span class="mi">10</span>
<span class="n">xb</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">rand</span><span class="p">(</span><span class="n">bs</span><span class="p">,</span> <span class="n">n_vars</span><span class="p">,</span> <span class="n">seq_len</span><span class="p">)</span>
<span class="n">new_head</span> <span class="o">=</span> <span class="n">partial</span><span class="p">(</span><span class="n">conv_lin_3d_head</span><span class="p">,</span> <span class="n">d</span><span class="o">=</span><span class="p">(</span><span class="mi">5</span> <span class="p">,</span><span class="mi">2</span><span class="p">))</span>
<span class="n">net</span> <span class="o">=</span> <span class="n">MultiTSTPlus</span><span class="p">(</span><span class="n">n_vars</span><span class="p">,</span> <span class="n">c_out</span><span class="p">,</span> <span class="n">seq_len</span><span class="p">,</span> <span class="n">custom_head</span><span class="o">=</span><span class="n">new_head</span><span class="p">)</span>
<span class="nb">print</span><span class="p">(</span><span class="n">net</span><span class="o">.</span><span class="n">to</span><span class="p">(</span><span class="n">xb</span><span class="o">.</span><span class="n">device</span><span class="p">)(</span><span class="n">xb</span><span class="p">)</span><span class="o">.</span><span class="n">shape</span><span class="p">)</span>
<span class="n">net</span><span class="o">.</span><span class="n">head</span>
</pre></div>

    </div>
</div>
</div>

<div class="output_wrapper">
<div class="output">

<div class="output_area">

<div class="output_subarea output_stream output_stdout output_text">
<pre>torch.Size([8, 5, 2])
</pre>
</div>
</div>

<div class="output_area">



<div class="output_text output_subarea output_execute_result">
<pre>Sequential(
  (0): create_conv_lin_3d_head(
    (0): BatchNorm1d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
    (1): Conv1d(128, 5, kernel_size=(1,), stride=(1,), bias=False)
    (2): Transpose(-1, -2)
    (3): BatchNorm1d(12, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
    (4): Transpose(-1, -2)
    (5): Linear(in_features=12, out_features=2, bias=False)
  )
)</pre>
</div>

</div>

</div>
</div>

</div>
    {% endraw %}

</div>
 

